package com.zhiyou.app.common.ws;

import cn.hutool.core.util.NumberUtil;
import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.xinchao.components.redis.utils.RedisUtil;
import com.zhiyou.app.common.constant.OrderTaskCacheMapping;
import com.zhiyou.app.common.exception.BizException;
import com.zhiyou.app.common.starejob.StartJob;
import com.zhiyou.app.common.util.SseEmitterSender;
import com.zhiyou.app.pojo.entity.TaskEntity;
import com.zhiyou.app.pojo.enums.TaskStatusEnum;
import com.zhiyou.app.service.ITaskService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.Arrays;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import static com.zhiyou.app.common.constant.CommonConstant.*;


@Slf4j
@Component
public class TaskProcessTracker implements StartJob {
    @Value("${comfyui.server}")
    private String serverUrl;
    @Resource
    private TaskFinishHandler taskFinishHandler;
    @Resource
    private ITaskService taskService;
    @Resource
    private RedisUtil redisUtil;
    @Resource
    private OrderTaskCacheMapping orderTaskCacheMapping;

    @Override
    public void run() throws Exception {
        CURRENT_TASK_INFO.setChangeListener((t) -> {
            ComfyuiTaskInfo newValue = t.getNewValue();
            ComfyuiTaskInfo oldValue = t.getOldValue();
            if (oldValue == null) {
                oldValue = new ComfyuiTaskInfo();
            }
            log.info("{}任务状态变化：{}->{}", newValue.getTask_id(), oldValue.getStatus(), newValue.getStatus());
            try {
                log.info("检测进度状态");
                //老任务完成
                if (!newValue.getTask_id().equals(oldValue.getTask_id()) && !redisUtil.hasKey(String.format(REDIS_OLD_TASK_COMPLETE_FLAG_PREFIX, oldValue.getTask_id()))) {
                    if (oldValue.getTask_id() != null) {
                        log.info("{}任务【old】完成, 速度太快未监听到complete，手动执行任务完成", oldValue.getTask_id());
                        finishTask(orderTaskCacheMapping.getOrderNumByPid(oldValue.getTask_id()), oldValue);
                        orderTaskCacheMapping.removeByPid(oldValue.getTask_id());
                    } else {
                        return;
                    }
                }
                switch (newValue.getStatus()) {
                    case "running":
                        if (!newValue.getTask_id().equals(oldValue.getTask_id()) && redisUtil.hasKey(String.format(REDIS_TASK_ID_PREFIX, newValue.getTask_id()))) {
                            log.info("{}任务开始执行", newValue.getTask_id());
                            SseEmitterSender.send(orderTaskCacheMapping.getOrderNumByPid(newValue.getTask_id()), "[state]:任务开始执行");
                        }
                        calcProcess(newValue, orderTaskCacheMapping.getOrderNumByPid(newValue.getTask_id()));
                        break;
                    case "error":
                    case "completed":
                        log.info("{}任务完成", newValue.getTask_id());
                        redisUtil.setEx(String.format(REDIS_OLD_TASK_COMPLETE_FLAG_PREFIX, newValue.getTask_id()), "1", 1, TimeUnit.DAYS);
                        finishTask(orderTaskCacheMapping.getOrderNumByPid(newValue.getTask_id()), newValue);
                        orderTaskCacheMapping.removeByPid(newValue.getTask_id());
                        break;
                }
            } catch (IOException e) {
                log.error("进度获取失败", e);
                throw new BizException("进度获取失败");
            }
        });
        Executors.newSingleThreadScheduledExecutor().scheduleWithFixedDelay(this::doGetTaskInfo, 5, 1, TimeUnit.SECONDS);
        //每分钟检测任务状态
        Executors.newSingleThreadScheduledExecutor().scheduleAtFixedRate(() -> {
            log.info("整体检测任务状态");
            taskService.timerGetTaskResult();
        }, 5, 60, TimeUnit.SECONDS);
    }

    private void finishTask(String orderNum, ComfyuiTaskInfo taskValue) throws IOException {
        boolean error = "error".equals(taskValue.getStatus());
        log.info("{}任务执行结束" + (error  ? "[异常]" : "[正常]"), taskValue.getTask_id());
        TaskEntity task = taskService.getOne(new LambdaQueryWrapper<TaskEntity>().eq(TaskEntity::getTaskPid, taskValue.getTask_id()));
        if (task != null && task.getUserId() != null) {
            taskFinishHandler.handle(taskValue.getTask_id(), orderNum, task, taskValue);
        }
        if(!error){
            calcProcess(taskValue, orderNum);
            SseEmitterSender.send(orderNum, "[state]:任务执行结束[" + taskValue.getTask_id() + "]");
            if (taskService.count(new LambdaQueryWrapper<TaskEntity>().eq(TaskEntity::getOrderNum, orderNum).in(TaskEntity::getTaskStatus, Arrays.asList(TaskStatusEnum.PROCESSING, TaskStatusEnum.CREATE))) == 0) {
                log.info("{}任务全部完成", orderNum);
                SseEmitterSender.sendWithDone(orderNum, "[state]:订单任务全部完成");
            }
        } else {
            SseEmitterSender.sendWithDone(orderNum, "[state]:任务异常");
        }
    }

    private void calcProcess(ComfyuiTaskInfo newValue, String sseKey) throws IOException {
        ComfyuiTaskInfo.WorkflowProgress workflowProgress = newValue.getWorkflow_progress();
        Double hasRun = Double.parseDouble(workflowProgress.getExecuted_nodes() + "");
        Double total = Double.parseDouble(workflowProgress.getTotal_nodes() + "");
        float totalPercent = NumberUtil.round(hasRun / total * 100, 2).floatValue();
        totalPercent = Math.min(totalPercent, 100f);
        SseEmitterSender.send(sseKey, "[process]:" + totalPercent);
        log.info("任务{} process:{}%", newValue.getTask_id(), totalPercent);
    }

    private void doGetTaskInfo() {
        HttpResponse response = HttpUtil.createGet(serverUrl.substring(0, serverUrl.length() - 5) + "/task_monitor/status").execute();
        if (response.isOk()) {
            CURRENT_TASK_INFO.set(JSONUtil.toBean(response.body(), ComfyuiTaskInfo.class));
        } else {
            log.error("获取comfyui插件任务信息失败");
        }
    }

    public void getTaskProcess(String promptId, String orderNum, TaskFinishHandler taskFinishHandler) throws IOException {
        this.taskFinishHandler = taskFinishHandler;
        doGetTaskInfo();
        try {
            TimeUnit.SECONDS.sleep(2);
        } catch (InterruptedException e) {
            throw new BizException("获取任务信息失败");
        }
        if (CURRENT_TASK_INFO.get().getQueue().getPending().stream().anyMatch(taskItem -> taskItem.getPrompt_id().equals(promptId))) {
            SseEmitterSender.send(orderNum, "[state]:任务排队中");
        }
    }
}
